import json

import uvicorn
from fastapi import FastAPI, APIRouter, HTTPException
from fastapi.responses import StreamingResponse
from pydantic import BaseModel

from agent.GovInfoAccess import GovInfoAccess


class GovAIService:
    def __init__(self):
        self.app = FastAPI(title="政府信息服务API")
        self.router = APIRouter()
        self.ai_agent = GovInfoAccess("./prompts/gov_qa_prompt_template.txt")

        # 添加路由
        self.router.add_api_route("/chat", self.chat_endpoint, methods=["POST"])
        self.router.add_api_route("/history/{session_id}", self.get_history, methods=["GET"])

        self.app.include_router(self.router)

    # 新增请求响应模型
    class ChatRequest(BaseModel):
        input: str
        session_id: str = None


    async def chat_endpoint(self, request: ChatRequest):
        """流式对话接口"""

        def generate():
            # 保留原有chat方法的流处理逻辑
            response_generator = self.ai_agent.stream_chat(
                input=request.input,
                session_id=request.session_id
            )

            for chunk in response_generator:
                yield json.dumps({
                    "text": chunk
                }, ensure_ascii=False).encode('utf-8') + b"\n"

        return StreamingResponse(generate(), media_type="text/event-stream")

    async def get_history(self, session_id: str):
        """历史记录查询接口"""
        history = self.ai_agent._get_long_term_context(session_id)
        if not history:
            raise HTTPException(status_code=404, detail="Session not found")
        return {
            "session_id": session_id,
            "history": [{"query": h[0], "response": h[1]} for h in history]
        }

    def run(self, host: str = "0.0.0.0", port: int = 8000):
        """启动API服务"""
        uvicorn.run(self.app, host=host, port=port)


# 修改使用示例部分为API启动
if __name__ == "__main__":
    service = GovAIService()
    service.run()
